import numpy as np
import cv2
from ocr_rec import TextRecognizer, init_args
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
import warnings
import argparse
warnings.filterwarnings("ignore")
from typing import List, Dict, Tuple, Optional, Union

global plate # 声明全局变量plate，用于存储当前识别到的车牌号
plate = ''# 初始化全局变量plate为空字符串

class PlateRecognizer:
    """
    车牌识别类，用于检测和识别图像中的车牌
    """
    def __init__(self, det_model_path="license_models/y11n-pose_plate_best.onnx"):
        # 初始化车牌检测模型（YOLO）
        self.model_det = YOLO(det_model_path)# 加载YOLO目标检测模型（onnx格式）
        # 初始化文字识别模型参数和实例
        parser = init_args().parse_args()# 解析文字识别所需的参数
        self.model_ocr = TextRecognizer(parser) # 创建文字识别器实例

    def recognize_plate(self, img):
        plate_objs=[] # 存储识别到的车牌信息列表
        # detect plates
        # 检测图像中的车牌
        plates = self.model_det(img, verbose=False)# 调用YOLO模型检测车牌，关闭详细输出

        # 遍历检测到的车牌（boxes.xyxy是边界框，boxes.conf是置信度）
        for plate, conf in zip(plates[0].boxes.xyxy, plates[0].boxes.conf):
        # 将边界框坐标从张量转换为整数（x1,y1是左上角，x2,y2是右下角）
            x1, y1, x2, y2 = plate.cpu() # 将数据从GPU转移到CPU（如果使用GPU）
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # 转换为整数坐标
            # 从原图中裁剪出车牌区域
            plate_img = img[y1:y2, x1:x2]
            # recognize text

            # 对裁剪的车牌图像进行文字识别
            try:
                rec_res, _ = self.model_ocr([plate_img])# 调用文字识别模型
            except Exception as E:
                print(E)# 打印异常信息
                exit() # 异常时退出程序

            # 整理识别结果（仅保留有效结果）
            if len(rec_res[0])>0:
                obj = {}# 存储单个车牌的信息
                obj['text'] = rec_res[0][0]# 识别到的车牌文字
                obj['score_text'] = rec_res[0][1] # 文字识别的置信度
                obj['bbox'] = [x1, y1, x2, y2]# 车牌边界框坐标
                obj['score_bbox'] = conf.cpu().numpy().item()# 检测框的置信度（转换为Python数值）
                plate_objs.append(obj)# 将结果添加到列表
        return plate_objs # 返回所有识别到的车牌信息
      

# def DrawPlateNum(img, palte_num, x1, y1):
#     img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
#     img_pil = Image.fromarray(img_rgb)
#     draw = ImageDraw.Draw(img_pil)
#     font = ImageFont.truetype("simsun.ttc", 40)
#     draw.text((x1, y1), palte_num, font=font, fill=(255, 255, 0))
#     img_bgr = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
#     return img_bgr

def process_video(video_path=None, output_path='output.mp4', frame_skip=0):
    # 
    # 处理视频文件或摄像头输入
    # :param video_path: 视频文件路径，如果为None则使用摄像头
    # # :param output_path: 输出视频文件路径
    # :param frame_skip: 跳帧数，例如设置为2表示每处理1帧跳过2帧
    # 
    # 初始化视频捕获
    if video_path is None:
        cap = cv2.VideoCapture(0)  # 使用默认摄像头
    else:
        cap = cv2.VideoCapture(video_path) # 从视频文件读取
    
    # 检查视频源是否打开成功
    if not cap.isOpened():
        print("Error: Could not open video source")
        return
    
    # 获取视频属性（帧率、宽度、高度）
    fps = cap.get(cv2.CAP_PROP_FPS)# 帧率
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))# 宽度
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# 高度
    
    # # 创建视频写入对象
    # fourcc = cv2.VideoWriter_fourcc(*'mp4v')# 视频编码格式
    # out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) # 初始化写入器
    
    # 初始化车牌识别器
    plate_rec = PlateRecognizer() # 创建PlateRecognizer实例
    
    print(f"开始处理视频，跳帧数: {frame_skip}，按'q'键退出...")
    
    frame_count = 0# 帧计数器
    while True:
        ret, frame = cap.read()# 读取一帧图像（ret为是否成功，frame为图像数据）
        if not ret:
            break# 读取失败或结束时退出循环
            
        # 跳帧处理
        # frame_count += 1
        # if frame_skip > 0 and frame_count % (frame_skip + 1) != 0:
        #     # 写入原始帧（不处理）
        #     # if output_path:
        #     #     out.write(frame)
        #     # # 显示原始帧
        #     # cv2.imshow('License Plate Recognition', frame)
        #     # 检查是否退出
        #     if cv2.waitKey(1) & 0xFF == ord('q'):
        #         break
        #     continue
            
        # 执行车牌识别
        plate_objs = plate_rec.recognize_plate(frame)# 调用识别方法处理当前帧

        #写入全局变量
        # 更新全局变量plate（仅保留高置信度结果）
        global plate
        if plate_objs and plate_objs[0]['score_text'] > 0.8: # 置信度大于0.8才更新
            plate = plate_objs[0]['text']
        else:
            plate = None # 低置信度时设为None
        
        # 在帧上绘制结果
        for bbox in plate_objs:
            if bbox['score_text'] > 0.8:  # 检查每个检测结果的置信度# 仅绘制高置信度结果
                x1, y1, x2, y2 = bbox['bbox'] # 获取边界框坐标
                # 绘制边界框（绿色，线宽2）
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                # 转换图像格式（BGR→RGB），用于PIL绘制中
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img_pil = Image.fromarray(frame)# 转换为PIL图像
                draw = ImageDraw.Draw(img_pil)# 创建绘制对象

                 # 加载中文字体（宋体，大小30）
                font = ImageFont.truetype("simsun.ttc", 30)
                 # 在边界框上方绘制车牌文字和置信度（红色
                draw.text((x1, y1 - 30), f"{bbox['text']} {bbox['score_text']:.2f}", 
                        font=font, fill=(255, 0, 0))
                # 转换回OpenCV格式（RGB→BGR）
                frame = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
        
        # 显示处理后的帧
        cv2.imshow('License Plate Recognition', frame)
        
        # # 写入输出视频
        # out.write(frame)
        
        # 按'q'键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    # 释放资源
    cap.release()# 释放视频捕获对象
    out.release() # 释放视频写入对象（注释：原代码未启用）
    cv2.destroyAllWindows() # 关闭所有OpenCV窗口
    # print(f"处理完成，结果已保存到 {output_path}")
    print("处理完成") # 打印处理完成信息

def get_plate():
    # 获取当前的车牌号
    global plate# 声明使用全局变量plate
    return plate # 返回全局变量plate的值

# 注释：测试函数（未启用）
# def test():
#     global plate  # 声明使用全局变量 plate
#     index = 1
#     plate = plate + str(index)  # 将数字转换为字符串后再拼接

def demo():
    parser = init_args().parse_args()# 解析命令行参数
    
    # 获取跳帧数，默认为0（不跳帧）
    frame_skip = getattr(parser, 'frame_skip', 0)# 从参数中获取frame_skip，无则默认0
    
    # 根据参数判断处理类型：视频文件、单张图片或摄像头
    if hasattr(parser, 'video_path') and parser.video_path:
        # 处理视频文件
        process_video(video_path=parser.video_path, frame_skip=frame_skip)
    elif hasattr(parser, 'image_path') and parser.image_path:
        # 处理单张图片（保留原有功能）
        img = cv2.imread(parser.image_path)# 读取图片
        if img is None:
            print("Error: no image found")# 图片读取失败提示
            return
        plate_rec = PlateRecognizer()# 创建车牌识别器实例
        plate_objs = plate_rec.recognize_plate(img) # 识别图片中的车牌

        # 处理识别结果（仅保留高置信度结果）
        if plate_objs and plate_objs[0]['score_text'] > 0.8:
            print(plate_objs)# 打印识别结果
             # 更新全局变量plate
            plate = plate_objs[0]['text'] if plate_objs and plate_objs[0]['score_text'] > 0.8 else None

            # 注释：原代码保留了图像绘制功能，但实际未启用
            # # 在图片上绘制结果
            # for bbox in plate_objs:
            #     if bbox['score_text'] > 0.8:  # 检查每个检测结果的置信度
            #         x1, y1, x2, y2 = bbox['bbox']
            #         cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            #         img = DrawPlateNum(img, f"{bbox['text']} {bbox['score_text']:.2f}", x1, y1-30)
            
            # 显示并保存结果
            # cv2.imshow("License Plate Recognition", img)
            # cv2.waitKey(0)
            # cv2.imwrite("result.jpg", img)
        else:
            print("未检测到置信度大于0.8的车牌")
    else:
        # 如果没有指定文件，则使用摄像头
        process_video(frame_skip=frame_skip)

# 计算平均识别时间
def test_ocr_speed():
    plate_rec = PlateRecognizer()# 创建车牌识别器实例
    # 测试图片列表（可替换为实际测试图片路径）
    imgs = [cv2.imread("test_plate/ec4a268fc_r.jpg"),
     cv2.imread("test_plate/hcqccip6335913.jpg"),
     cv2.imread("test_plate/u=1570203173,1434495175&fm=253&fmt=auto&app=138&f=JPEG.webp"),
     ]
    #warm up
    # 预热模型（避免首次推理耗时影响测试）
    plate_objs = plate_rec.recognize_plate(imgs[0])
    # 开始计时（使用OpenCV的高精度计时器）
    time_start = cv2.getTickCount()
    cal_count = 0# 推理次数计数器
    while True: # 循环测试（无限循环，需手动终止）
        # 循环使用测试图片
        plate_objs = plate_rec.recognize_plate(imgs[cal_count%len(imgs)])
        cal_count += 1 # 累加次数
         # 计算耗时
        time_pause = cv2.getTickCount()
        time_cost = (time_pause - time_start)*1000 / cv2.getTickFrequency()# 转换为毫秒
         # 打印平均耗时和识别结果
        print(f"avgtime cost: {time_cost/(cal_count)}ms {plate_objs}")

if __name__ == "__main__":
    # 初始化参数解析器
    parser = init_args().parse_args()
     # 如果指定了speed_test参数，则运行速度测试
    if parser.speed_test:
        test_ocr_speed()
    else:
      # 否则运行演示功能        
        demo()

